#  导入需要的库
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder, StandardScaler
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error
from sklearn.pipeline import Pipeline  # 导入Pipeline
import numpy as np

# 加载数据
df = pd.read_excel('D:\PyCharm 2022.3\douban\豆瓣电影top250数据.xlsx')

# 数据预处理
def extract_country_and_genre(row):
    info = row['上映时间、电影类型'].split('/')
    country = info[0].split(' ')[-1]  # 提取国家
    genres = set(info[1:])  # 提取电影类型，并去重
    return country, genres
df[['country', 'genres']] = df.apply(lambda row: pd.Series(extract_country_and_genre(row)), axis=1).apply(pd.Series)

# 对国家进行编码
label_encoder = LabelEncoder()
df['country_encoded'] = label_encoder.fit_transform(df['country'])
all_genres = ['动作', '喜剧', '爱情', '科幻', '悬疑', '惊悚', '犯罪', '冒险', '动画', '奇幻', '家庭', '战争',
              '音乐', '历史', '传记', '纪录片', '运动', '儿童', '剧情', '恐怖', '灾难', '古装', '奇幻冒险', '爱情喜剧', '犯罪悬疑']

# 创建固定长度的类型特征向量
def create_genre_vector(genres, all_genres):
    genre_vector = [0] * len(all_genres)
    for genre in genres:
        if genre in all_genres:
            genre_vector[all_genres.index(genre)] = 1
    return genre_vector
df['genre_vector'] = df['genres'].apply(lambda genres: create_genre_vector(genres, all_genres))

# 准备特征和标签
X = np.column_stack((df[['country_encoded']].values, np.array(df['genre_vector'].tolist())))  # 这里假设你已经将genre_vector列转换为了列表的列表（每部电影一个列表）
y = df['电影评分']

# 拆分为训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 构建回归模型管道
pipeline = Pipeline([
    ('scaler', StandardScaler()),  # 标准化数据
    ('regressor', RandomForestRegressor(n_estimators=100, random_state=42))  # 使用随机森林回归器
])

# 训练模型
pipeline.fit(X_train, y_train)

# 预测并评估模型
y_pred = pipeline.predict(X_test)
mse = mean_squared_error(y_test, y_pred)
print(f'均方误差: {mse}')


def predict_movie_rating(new_sample_info):
    # 从新样本信息中提取国家和电影类型
    country, genres = extract_country_and_genre({"上映时间、电影类型": new_sample_info})
    # 对国家进行编码（使用训练过程中已拟合的LabelEncoder）
    country_encoded = label_encoder.transform([country])[0]
    # 创建电影类型的特征向量
    genre_vector = create_genre_vector(genres, all_genres)
    # 将新样本的特征组合成模型所需的输入格式
    new_sample_array = np.array([[country_encoded] + genre_vector]).astype(float)
    # 标准化新样本的特征（使用训练过程中已拟合的Scaler）
    new_sample_array_scaled = pipeline.named_steps['scaler'].transform(new_sample_array)
    # 使用训练好的模型进行预测
    predicted_rating = pipeline.predict(new_sample_array_scaled)
    return predicted_rating[0]

# 示例新样本信息
new_sample = "2023年 美国 / 动作 / 科幻"

# 预测新样本的电影评分
predicted_rating = predict_movie_rating(new_sample)
print(f"预测的电影评分为: {predicted_rating}")
